{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "e45f9b60-cd6b-4c15-958f-1feca5438128",
   "metadata": {},
   "source": [
    "# SQL Index Guide (Core)\n",
    "\n",
    "This is a basic guide to LlamaIndex's SQL index capabilities. We first show how to define a SQL table, then we build a TableIndex over the schema. This will allow us to synthesize a SQL query given the user's natural language query."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "119eb42b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# import logging\n",
    "# import sys\n",
    "\n",
    "# logging.basicConfig(stream=sys.stdout, level=logging.INFO)\n",
    "# logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "107396a9-4aa7-49b3-9f0f-a755726c19ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import Markdown, display"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "461438c8-302d-45c5-8e69-16ad604686d1",
   "metadata": {},
   "source": [
    "### Create Database Schema\n",
    "\n",
    "We use `sqlalchemy`, a popular SQL database toolkit, to create an empty `city_stats` Table"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "a370b266-66f5-4624-bbf9-2ad57f0511f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sqlalchemy import (\n",
    "    create_engine,\n",
    "    MetaData,\n",
    "    Table,\n",
    "    Column,\n",
    "    String,\n",
    "    Integer,\n",
    "    select,\n",
    "    column,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "ea24f794-f10b-42e6-922d-9258b7167405",
   "metadata": {},
   "outputs": [],
   "source": [
    "engine = create_engine(\"sqlite:///:memory:\")\n",
    "metadata_obj = MetaData()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "b4154b29-7e23-4c26-a507-370a66186ae7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# create city SQL table\n",
    "table_name = \"city_stats\"\n",
    "city_stats_table = Table(\n",
    "    table_name,\n",
    "    metadata_obj,\n",
    "    Column(\"city_name\", String(16), primary_key=True),\n",
    "    Column(\"population\", Integer),\n",
    "    Column(\"country\", String(16), nullable=False),\n",
    ")\n",
    "metadata_obj.create_all(engine)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "1c09089a-6bcd-48db-8120-a84c8da3f82e",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Define SQL Database\n",
    "\n",
    "We first define our `SQLDatabase` abstraction (a light wrapper around SQLAlchemy). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "768d1581-b482-4c73-9963-5ffd68a2aafb",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from llama_index import SQLDatabase, ServiceContext\n",
    "from llama_index.llms import OpenAI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "bffabba0-8e54-4f24-ad14-2c8979c582a5",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "llm = OpenAI(temperature=0, model=\"text-davinci-002\")\n",
    "service_context = ServiceContext.from_defaults(llm=llm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "9432787b-a8f0-4fc3-8323-e2cd9497df73",
   "metadata": {},
   "outputs": [],
   "source": [
    "sql_database = SQLDatabase(engine, include_tables=[\"city_stats\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "84d4ee54-9f00-40fd-bab0-36e5e579dc9f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'\\nCREATE TABLE city_stats (\\n\\tcity_name VARCHAR(16) NOT NULL, \\n\\tpopulation INTEGER, \\n\\tcountry VARCHAR(16) NOT NULL, \\n\\tPRIMARY KEY (city_name)\\n)\\n\\n/*\\n3 rows from city_stats table:\\ncity_name\\tpopulation\\tcountry\\n\\n*/'"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sql_database.table_info"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "bad7ffbe",
   "metadata": {},
   "source": [
    "We add some testing data to our SQL database."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "95043e10-6cdf-4f66-96bd-ce307ea7df3e",
   "metadata": {},
   "outputs": [],
   "source": [
    "sql_database = SQLDatabase(engine, include_tables=[\"city_stats\"])\n",
    "from sqlalchemy import insert\n",
    "\n",
    "rows = [\n",
    "    {\"city_name\": \"Toronto\", \"population\": 2930000, \"country\": \"Canada\"},\n",
    "    {\"city_name\": \"Tokyo\", \"population\": 13960000, \"country\": \"Japan\"},\n",
    "    {\"city_name\": \"Chicago\", \"population\": 2679000, \"country\": \"United States\"},\n",
    "    {\"city_name\": \"Seoul\", \"population\": 9776000, \"country\": \"South Korea\"},\n",
    "]\n",
    "for row in rows:\n",
    "    stmt = insert(city_stats_table).values(**row)\n",
    "    with engine.connect() as connection:\n",
    "        cursor = connection.execute(stmt)\n",
    "        connection.commit()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "b315b8ff-7dd7-4e7d-ac47-8c5a0c3e7ae9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[('Toronto', 2930000, 'Canada'), ('Tokyo', 13960000, 'Japan'), ('Chicago', 2679000, 'United States'), ('Seoul', 9776000, 'South Korea')]\n"
     ]
    }
   ],
   "source": [
    "# view current table\n",
    "stmt = select(city_stats_table.c[\"city_name\", \"population\", \"country\"]).select_from(\n",
    "    city_stats_table\n",
    ")\n",
    "\n",
    "with engine.connect() as connection:\n",
    "    results = connection.execute(stmt).fetchall()\n",
    "    print(results)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "051a171f-8c97-40ed-ae17-4e3fa3785487",
   "metadata": {},
   "source": [
    "### Query Index"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "f6a2303f-3bae-4fa2-8750-03f9af747848",
   "metadata": {},
   "source": [
    "We first show how we can execute a raw SQL query, which directly executes over the table."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "eddd3608-31ff-4591-a02a-90987e312669",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "('Chicago',)\n",
      "('Seoul',)\n",
      "('Tokyo',)\n",
      "('Toronto',)\n"
     ]
    }
   ],
   "source": [
    "from sqlalchemy import text\n",
    "\n",
    "with engine.connect() as con:\n",
    "    rows = con.execute(text(\"SELECT city_name from city_stats\"))\n",
    "    for row in rows:\n",
    "        print(row)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "4e72b931",
   "metadata": {},
   "source": [
    "## Natural language SQL\n",
    "Once we have constructed our SQL database, we can use the NLSQLTableQueryEngine to\n",
    "construct natural language queries that are synthesized into SQL queries.\n",
    "\n",
    "Note that we need to specify the tables we want to use with this query engine.\n",
    "If we don't the query engine will pull all the schema context, which could\n",
    "overflow the context window of the LLM."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "5d992fb5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.indices.struct_store.sql_query import NLSQLTableQueryEngine\n",
    "\n",
    "query_engine = NLSQLTableQueryEngine(\n",
    "    sql_database=sql_database,\n",
    "    tables=[\"city_stats\"],\n",
    ")\n",
    "query_str = \"Which city has the highest population?\"\n",
    "response = query_engine.query(query_str)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "298b4ca2",
   "metadata": {},
   "source": [
    "This query engine should used in any case where you can specify the tables you want\n",
    "to query over beforehand, or the total size of all the table schema plus the rest of\n",
    "the prompt fits your context window."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "dee4d251",
   "metadata": {},
   "source": [
    "## Building our Table Index\n",
    "If we don't know ahead of time which table we would like to use, and the total size of\n",
    "the table schema overflows your context window size, we should store the table schema \n",
    "in an index so that during query time we can retrieve the right schema.\n",
    "\n",
    "The way we can do this is using the SQLTableNodeMapping object, which takes in a \n",
    "SQLDatabase and produces a Node object for each SQLTableSchema object passed \n",
    "into the ObjectIndex constructor.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "d71045c0-7a96-4e86-b38c-c378b7759aa4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.indices.struct_store.sql_query import SQLTableRetrieverQueryEngine\n",
    "from llama_index.objects import SQLTableNodeMapping, ObjectIndex, SQLTableSchema\n",
    "from llama_index import VectorStoreIndex\n",
    "\n",
    "# set Logging to DEBUG for more detailed outputs\n",
    "table_node_mapping = SQLTableNodeMapping(sql_database)\n",
    "table_schema_objs = [\n",
    "    (SQLTableSchema(table_name=\"city_stats\"))\n",
    "]  # add a SQLTableSchema for each table\n",
    "\n",
    "obj_index = ObjectIndex.from_objects(\n",
    "    table_schema_objs,\n",
    "    table_node_mapping,\n",
    "    VectorStoreIndex,\n",
    ")\n",
    "query_engine = SQLTableRetrieverQueryEngine(\n",
    "    sql_database, obj_index.as_retriever(similarity_top_k=1)\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "b6156caf",
   "metadata": {},
   "source": [
    "Now we can take our SQLTableRetrieverQueryEngine and query it for our response."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "802da9ed",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "<b> Tokyo has the highest population, with 13,960,000 people.</b>"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "response = query_engine.query(\"Which city has the highest population?\")\n",
    "display(Markdown(f\"<b>{response}</b>\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "54a99cb0-578a-40ec-a3eb-1666ac18fbed",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[('Tokyo', 13960000)]"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# you can also fetch the raw result from SQLAlchemy!\n",
    "response.metadata[\"result\"]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "0d19b9cd",
   "metadata": {},
   "source": [
    "You can also add additional context information for each table schema you define."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "44a87651",
   "metadata": {},
   "outputs": [],
   "source": [
    "# manually set context text\n",
    "city_stats_text = (\n",
    "    \"This table gives information regarding the population and country of a given city.\\n\"\n",
    "    \"The user will query with codewords, where 'foo' corresponds to population and 'bar'\"\n",
    "    \"corresponds to city.\"\n",
    ")\n",
    "\n",
    "table_node_mapping = SQLTableNodeMapping(sql_database)\n",
    "table_schema_objs = [\n",
    "    (SQLTableSchema(table_name=\"city_stats\", context_str=city_stats_text))\n",
    "]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "2a567566-9828-4aa4-afde-b253c01eda69",
   "metadata": {},
   "source": [
    "### Using LangChain for Querying\n",
    "\n",
    "Since our SQLDatabase inherits from langchain, you can also use langchain itself for querying purposes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "8e0acde4-ca61-42e9-97f8-c9cf11502157",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain import OpenAI, SQLDatabase, SQLDatabaseChain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "30860e8b-9ad0-418c-b266-753242c1f208",
   "metadata": {},
   "outputs": [],
   "source": [
    "llm = OpenAI(temperature=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "07068a3a-30a4-4473-ba82-ab6e93e3437c",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/hongyishi/Documents/GitHub/gpt_index/.venv/lib/python3.11/site-packages/langchain/chains/sql_database/base.py:63: UserWarning: Directly instantiating an SQLDatabaseChain with an llm is deprecated. Please instantiate with llm_chain argument or using the from_llm class method.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "# set Logging to DEBUG for more detailed outputs\n",
    "db_chain = SQLDatabaseChain(llm=llm, database=sql_database)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "a04c0a1d-f6a8-4a4a-9181-4123b09ec614",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Tokyo has the highest population with 13960000 people.'"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "db_chain.run(\"Which city has the highest population?\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
